"""KPS module"""
import sys
import mindspore as ms
from mindspore import nn
from mindspore import ops
import mind3d.models.blocks.groupfree_3d_pointnet2_util as groupfree_3d_pointnet2_util


class PointsObjClsModule(nn.Cell):
    """
    Object candidate point prediction from seed point features.

    Args:
        seed_feature_dim(int): number of channels of seed point features.
    """

    def __init__(self, seed_feature_dim):
        super(PointsObjClsModule, self).__init__()
        self.in_dim = seed_feature_dim
        self.conv1 = nn.Conv1d(self.in_dim, self.in_dim, kernel_size=1, has_bias=True, weight_init="he_uniform", bias_init="normal")
        self.bn1 = nn.BatchNorm2d(self.in_dim)
        self.conv2 = nn.Conv1d(self.in_dim, self.in_dim, kernel_size=1, has_bias=True, weight_init="he_uniform", bias_init="normal")
        self.bn2 = nn.BatchNorm2d(self.in_dim)
        self.conv3 = nn.Conv1d(self.in_dim, 1, kernel_size=1, has_bias=True, weight_init="he_uniform", bias_init="normal")
        self.relu = ops.ReLU()
        self.expand_dims = ops.ExpandDims()
        self.squeeze = ops.Squeeze(-1)

    def construct(self, seed_features):
        """
        Args:
            seed_features: (batch_size, feature_dim, num_seed) tensor
        Returns:
            logits: (batch_size, 1, num_seed)
        """
        net = self.relu(self.squeeze(self.bn1(self.expand_dims(self.conv1(seed_features), -1))))
        net = self.relu(self.squeeze(self.bn2(self.expand_dims(self.conv2(net), -1))))
        logits = self.conv3(net)  # (batch_size, 1, num_seed)
        return logits


class KPSModule(nn.Cell):
    """
    Args:
        num_proposal: int (default: 128)
                Number of proposals/detections generated from the network.
                Each proposal is a 3D OBB with a semantic class.
        seed_feature_dim: int (default: 288)
                number of channels of seed point features
                
                
    Input:
    xyz: (B,N,3)
    features: (B,C,N)
    Output:
    new_xyz: (B,P,3)
    new_features: (B,C,P)
    
    """

    def __init__(self, seed_feature_dim=288, num_proposal=256):
        super(KPSModule, self).__init__()
        self.seed_feature_dim = seed_feature_dim
        self.num_proposal = num_proposal
        self.points_obj_cls = PointsObjClsModule(self.seed_feature_dim)

    def construct(self, xyz, features):
        transpose = ops.Transpose()
        points_obj_cls_logits = self.points_obj_cls(features) #(B,1,N)
        points_obj_cls_scores = ops.Sigmoid()(points_obj_cls_logits)
        points_obj_cls_scores = ops.Squeeze(1)(points_obj_cls_scores)
        sample_inds = ops.Cast()( ops.TopK(sorted=True)(points_obj_cls_scores, self.num_proposal)[1], ms.int32)

        new_xyz = groupfree_3d_pointnet2_util.index_points(xyz, sample_inds)
        new_features = transpose(groupfree_3d_pointnet2_util.index_points(transpose(features, (0, 2, 1)), sample_inds), (0, 2, 1))

        return new_xyz, new_features, sample_inds, points_obj_cls_logits
